{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 模型蒸馏-白盒蒸馏",
   "id": "6b61d271e60ad83e"
  },
  {
   "cell_type": "code",
   "id": "82e9431d0c343a30",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-02T23:00:27.543903Z",
     "start_time": "2025-08-02T23:00:12.362340Z"
    }
   },
   "source": [
    "from peft import PeftModel\n",
    "\n",
    "from modelscope import AutoModelForCausalLM\n",
    "from modelscope import AutoTokenizer\n",
    "\n",
    "teacher_model_name = \"Qwen/Qwen3-1.7B\"\n",
    "teacher_base_model = AutoModelForCausalLM.from_pretrained(teacher_model_name)\n",
    "teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_model_name)\n",
    "teacher_model = PeftModel.from_pretrained(teacher_base_model, \"./Zhouyu-Qwen3-1.7B\")"
   ],
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/dadudu/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading Model from https://www.modelscope.cn to directory: /Users/dadudu/.cache/modelscope/hub/models/Qwen/Qwen3-1.7B\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-03 07:00:16,449 - modelscope - INFO - Target directory already exists, skipping creation.\n",
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:09<00:00,  4.56s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading Model from https://www.modelscope.cn to directory: /Users/dadudu/.cache/modelscope/hub/models/Qwen/Qwen3-1.7B\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-03 07:00:27,082 - modelscope - INFO - Target directory already exists, skipping creation.\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-02T23:00:44.972043Z",
     "start_time": "2025-08-02T23:00:38.172432Z"
    }
   },
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>user\n",
      "你是谁<|im_end|>\n",
      "<|im_start|>assistant\n",
      "<think>\n",
      "\n",
      "</think>\n",
      "\n",
      "我是大都督周瑜的AI助手\n",
      "<|im_end|>\n"
     ]
    }
   ],
   "execution_count": 2,
   "source": [
    "prompt = \"你是谁\"\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": prompt}\n",
    "]\n",
    "text = teacher_tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=True,\n",
    "    enable_thinking=False\n",
    ")\n",
    "model_inputs = teacher_tokenizer([text], return_tensors=\"pt\")\n",
    "generated_ids = teacher_model.generate(**model_inputs, max_new_tokens=256)\n",
    "content = teacher_tokenizer.decode(generated_ids[0])\n",
    "print(content)"
   ],
   "id": "e4fda820-fe36-4458-8eaa-558be2a65594"
  },
  {
   "cell_type": "code",
   "id": "8115e80e018fbe33",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-02T23:00:56.889145Z",
     "start_time": "2025-08-02T23:00:49.606904Z"
    }
   },
   "source": [
    "# Student Model\n",
    "from modelscope import AutoModelForCausalLM\n",
    "from modelscope import AutoTokenizer\n",
    "\n",
    "student_model_name = \"Qwen/Qwen3-0.6B\"\n",
    "student_model = AutoModelForCausalLM.from_pretrained(student_model_name)\n",
    "student_tokenizer = AutoTokenizer.from_pretrained(student_model_name)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading Model from https://www.modelscope.cn to directory: /Users/dadudu/.cache/modelscope/hub/models/Qwen/Qwen3-0.6B\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-03 07:00:50,788 - modelscope - INFO - Target directory already exists, skipping creation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading Model from https://www.modelscope.cn to directory: /Users/dadudu/.cache/modelscope/hub/models/Qwen/Qwen3-0.6B\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-03 07:00:56,730 - modelscope - INFO - Target directory already exists, skipping creation.\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "id": "9a248feea6ee272b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-02T23:01:02.906296Z",
     "start_time": "2025-08-02T23:01:02.888416Z"
    }
   },
   "source": [
    "# 蒸馏数据\n",
    "distillation_data = [\"你是谁\"]\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "chat_prompt_template = \"\"\"\n",
    "<|im_start|>user\n",
    "{question}<|im_end|>\n",
    "<|im_start|>assistant\n",
    "<think>\n",
    "\n",
    "</think>\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class ZhouyuDistillationDataset(Dataset):\n",
    "    def __init__(self, data, max_length=128):\n",
    "        self.encodings = []\n",
    "        for question in data:\n",
    "            text = chat_prompt_template.format(question=question)\n",
    "            print(text)\n",
    "            encoded = teacher_tokenizer(\n",
    "                [text],\n",
    "                return_tensors='pt'\n",
    "            )\n",
    "            input_ids = encoded['input_ids'].squeeze()\n",
    "            self.encodings.append(input_ids)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.encodings)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.encodings[idx]\n",
    "\n",
    "\n",
    "distillation_dataset = ZhouyuDistillationDataset(distillation_data)\n",
    "distillation_dataloader = DataLoader(distillation_dataset, batch_size=1, shuffle=True)\n",
    "for batch in distillation_dataloader:\n",
    "    print(batch)\n",
    "    break"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "<|im_start|>user\n",
      "你是谁<|im_end|>\n",
      "<|im_start|>assistant\n",
      "<think>\n",
      "\n",
      "</think>\n",
      "\n",
      "tensor([[   198, 151644,    872,    198, 105043, 100165, 151645,    198, 151644,\n",
      "          77091,    198, 151667,    271, 151668,    198]])\n"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "id": "99a1b365e61f938f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-02T23:13:09.866486Z",
     "start_time": "2025-08-02T23:07:46.110935Z"
    }
   },
   "source": [
    "import torch\n",
    "\n",
    "optimizer = torch.optim.Adam(student_model.parameters(), lr=5e-5)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "teacher_model.to(device)\n",
    "student_model.to(device)\n",
    "\n",
    "EPOCHS = 50\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "\n",
    "    for input_ids in distillation_dataloader:\n",
    "        input_ids = input_ids.to(device)\n",
    "\n",
    "        # 得到下一个词的概率分布\n",
    "        # with torch.no_grad():\n",
    "        #     teacher_outputs = []\n",
    "        #     current_inputs = input_ids.clone()\n",
    "        #\n",
    "        #     # 自回归生成\n",
    "        #     for _ in range(input_ids.size(1)):\n",
    "        #         outputs = teacher_model(current_inputs)\n",
    "        #         teacher_logits = outputs.logits[:, -1, :]\n",
    "        #         teacher_outputs.append(teacher_logits)\n",
    "        #\n",
    "        #         # 取概率最大的token继续生成\n",
    "        #         next_tokens = torch.argmax(teacher_logits, dim=-1)\n",
    "        #         current_inputs = torch.cat([current_inputs, next_tokens.unsqueeze(-1)], dim=-1)\n",
    "        #\n",
    "        #     teacher_logits = torch.stack(teacher_outputs, dim=1)  # [batch, seq_len, vocab]\n",
    "        #\n",
    "        # student_outputs = []\n",
    "        # current_inputs = input_ids.clone()\n",
    "        # for i in range(input_ids.size(1)):\n",
    "        #     outputs = student_model(current_inputs)\n",
    "        #     student_logits = outputs.logits[:, -1, :]\n",
    "        #     student_outputs.append(student_logits)\n",
    "        #\n",
    "        #     # 取概率最大的token继续生成\n",
    "        #     next_tokens = torch.argmax(student_logits, dim=-1)\n",
    "        #     current_inputs = torch.cat([current_inputs, next_tokens.unsqueeze(-1)], dim=-1)\n",
    "        #\n",
    "        # student_logits = torch.stack(student_outputs, dim=1)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            teacher_outputs = teacher_model(input_ids, labels=input_ids)\n",
    "            teacher_logits = teacher_outputs.logits\n",
    "\n",
    "        student_outputs = student_model(\n",
    "            input_ids, labels=input_ids\n",
    "        )\n",
    "        student_logits = student_outputs.logits\n",
    "\n",
    "        # 软目标损失\n",
    "        soft_loss = torch.nn.functional.kl_div(\n",
    "            torch.log_softmax(student_logits, dim=-1),\n",
    "            torch.softmax(teacher_logits, dim=-1),\n",
    "            reduction='batchmean'\n",
    "        )\n",
    "\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        soft_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if epoch % 10 == 0:\n",
    "            print(f'Epoch {epoch + 1}, Loss: {soft_loss:.4f}')\n",
    "\n"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CausalLMOutputWithPast(loss=tensor(1.1071), logits=tensor([[[ 4.1185,  5.5874,  7.1719,  ...,  0.4592,  0.4592,  0.4592],\n",
      "         [ 2.2171,  3.7658,  7.7227,  ..., -0.1603, -0.1604, -0.1604],\n",
      "         [ 2.7057,  1.1595,  4.4529,  ...,  0.3835,  0.3835,  0.3835],\n",
      "         ...,\n",
      "         [-8.8077, -3.2438, -4.8060,  ..., -0.8007, -0.8007, -0.8007],\n",
      "         [ 2.5121,  5.8145,  6.2906,  ..., -0.6621, -0.6621, -0.6620],\n",
      "         [-1.7819,  4.2719, -0.8884,  ..., -0.3974, -0.3974, -0.3976]]]), past_key_values=<transformers.cache_utils.DynamicCache object at 0x3d79b8550>, hidden_states=None, attentions=None)\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[8], line 45\u001B[0m\n\u001B[1;32m     42\u001B[0m combined_loss \u001B[38;5;241m=\u001B[39m SOFT_LOSS_WEIGHT \u001B[38;5;241m*\u001B[39m soft_loss \u001B[38;5;241m+\u001B[39m HARD_LOSS_WEIGHT \u001B[38;5;241m*\u001B[39m hard_loss\n\u001B[1;32m     44\u001B[0m optimizer\u001B[38;5;241m.\u001B[39mzero_grad()\n\u001B[0;32m---> 45\u001B[0m \u001B[43mcombined_loss\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbackward\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     46\u001B[0m optimizer\u001B[38;5;241m.\u001B[39mstep()\n\u001B[1;32m     48\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m epoch \u001B[38;5;241m%\u001B[39m \u001B[38;5;241m10\u001B[39m \u001B[38;5;241m==\u001B[39m \u001B[38;5;241m0\u001B[39m:\n",
      "File \u001B[0;32m~/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/torch/_tensor.py:648\u001B[0m, in \u001B[0;36mTensor.backward\u001B[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001B[0m\n\u001B[1;32m    638\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m has_torch_function_unary(\u001B[38;5;28mself\u001B[39m):\n\u001B[1;32m    639\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m handle_torch_function(\n\u001B[1;32m    640\u001B[0m         Tensor\u001B[38;5;241m.\u001B[39mbackward,\n\u001B[1;32m    641\u001B[0m         (\u001B[38;5;28mself\u001B[39m,),\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    646\u001B[0m         inputs\u001B[38;5;241m=\u001B[39minputs,\n\u001B[1;32m    647\u001B[0m     )\n\u001B[0;32m--> 648\u001B[0m \u001B[43mtorch\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mautograd\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbackward\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    649\u001B[0m \u001B[43m    \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mgradient\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mretain_graph\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcreate_graph\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43minputs\u001B[49m\n\u001B[1;32m    650\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/torch/autograd/__init__.py:307\u001B[0m, in \u001B[0;36mbackward\u001B[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001B[0m\n\u001B[1;32m    243\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[38;5;21mbackward\u001B[39m(\n\u001B[1;32m    244\u001B[0m     tensors: _TensorOrTensorsOrGradEdge,\n\u001B[1;32m    245\u001B[0m     grad_tensors: Optional[_TensorOrTensors] \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m,\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    249\u001B[0m     inputs: Optional[_TensorOrTensorsOrGradEdge] \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m,\n\u001B[1;32m    250\u001B[0m ) \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m>\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[1;32m    251\u001B[0m \u001B[38;5;250m    \u001B[39m\u001B[38;5;124mr\u001B[39m\u001B[38;5;124;03m\"\"\"Compute the sum of gradients of given tensors with respect to graph leaves.\u001B[39;00m\n\u001B[1;32m    252\u001B[0m \n\u001B[1;32m    253\u001B[0m \u001B[38;5;124;03m    The graph is differentiated using the chain rule. If any of ``tensors``\u001B[39;00m\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    305\u001B[0m \u001B[38;5;124;03m            were used to compute the :attr:`tensors`.\u001B[39;00m\n\u001B[1;32m    306\u001B[0m \u001B[38;5;124;03m    \"\"\"\u001B[39;00m\n\u001B[0;32m--> 307\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[43mtorch\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_C\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_are_functorch_transforms_active\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m:\n\u001B[1;32m    308\u001B[0m         \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mRuntimeError\u001B[39;00m(\n\u001B[1;32m    309\u001B[0m             \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mbackward() called inside a functorch transform. This is not \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    310\u001B[0m             \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124msupported, please use functorch.grad or functorch.vjp instead \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    311\u001B[0m             \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mor call backward() outside of functorch transforms.\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    312\u001B[0m         )\n\u001B[1;32m    314\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m grad_variables \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cbe0a8dd-5473-4407-afed-e6daec5ec9fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>user\n",
      "你是谁<|im_end|>\n",
      "<|im_start|>assistant\n",
      "<think>\n",
      "\n",
      "</think>\n",
      "\n",
      "我是大都督周瑜的AI助手\n",
      "<|im_end|>\n"
     ]
    }
   ],
   "source": [
    "prompt = \"你是谁\"\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": prompt}\n",
    "]\n",
    "text = student_tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=True,\n",
    "    enable_thinking=False,\n",
    ")\n",
    "model_inputs = student_tokenizer([text], return_tensors=\"pt\")\n",
    "model_inputs = model_inputs.to(device)\n",
    "generated_ids = student_model.generate(**model_inputs, max_new_tokens=256)\n",
    "content = student_tokenizer.decode(generated_ids[0])\n",
    "print(content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
